import numpy as np
from tqdm import tqdm
import torch
from torch.cuda.amp import autocast as autocast
from sklearn.metrics import confusion_matrix
from utils import save_imgs


def train_one_epoch(train_loader,
                    model,
                    criterion, 
                    optimizer, 
                    scheduler,
                    epoch, 
                    logger, 
                    config, 
                    scaler=None):
    '''
    train model for one epoch
    '''
    # switch to train mode
    model.train() 
 
    loss_list = []

    for iter, data in enumerate(train_loader):
        optimizer.zero_grad()
        images, targets = data
        # print(targets)
        # exit()
        images, targets = images.cuda(non_blocking=True).float(), targets.cuda(non_blocking=True).float()
        
        if config.amp:
            with autocast():
                out = model(images)
                loss = criterion(out, targets)      
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            out = model(images)
            loss = criterion(out, targets)
            loss.backward()
            optimizer.step()
        
        loss_list.append(loss.item())

        now_lr = optimizer.state_dict()['param_groups'][0]['lr']
        if iter % config.print_interval == 0:
            log_info = f'train: epoch {epoch}, iter:{iter}, loss: {np.mean(loss_list):.4f}, lr: {now_lr}'
            print(log_info)
            logger.info(log_info)
    scheduler.step() 


def val_one_epoch(test_loader,    
                  model,    
                  criterion,     
                  epoch,     
                  logger,    
                  config):    
    # 切换到评估模式    
    model.eval()    
    preds = []    
    gts = []    
    loss_list = []    
    with torch.no_grad():    
        for data in tqdm(test_loader):    
            img, msk = data    
            img, msk = img.cuda(non_blocking=True).float(), msk.cuda(non_blocking=True).float()    
            out = model(img)    
            loss = criterion(out, msk)    
            loss_list.append(loss.item())    
            # 假设msk已经是整数类型的标签（0和1），不需要进行任何处理    
            gts.append(msk.squeeze(1).cpu().numpy())    
            if type(out) is tuple:    
                out = out[0]  # 如果模型输出是元组，则取第一个元素作为预测    
            preds.append(out.squeeze(1).cpu().numpy())  # 假设out已经是概率分布，无需进一步处理    
    
    if epoch % config.val_interval == 0:    
        # 合并预测和真实标签    
        preds = np.concatenate(preds, axis=0)    
        gts = np.concatenate(gts, axis=0)    
    
        # 应用阈值将预测概率转换为二分类标签    
        y_pre = np.where(preds >= config.threshold, 1, 0).astype(np.uint8)    
    
        # 假设gts已经是整数类型的标签（0和1），无需进行任何处理    
        y_true = gts.astype(np.uint8)    
    
        # 计算混淆矩阵    
        confusion = confusion_matrix(y_true.ravel(), y_pre.ravel())    
    
        # 从混淆矩阵中提取TN, FP, FN, TP    
        TN = confusion[0, 0]    
        FP = confusion[0, 1]    
        FN = confusion[1, 0]    
        TP = confusion[1, 1]    
    
        # 这里可以继续进行其他性能指标的计算和日志记录    
        # ...    

        # 转换为浮点数
        TN = float(TN)
        FP = float(FP)
        FN = float(FN)
        TP = float(TP)

        accuracy = (TN + TP) / (TN + FP + FN + TP) if (TN + FP + FN + TP) != 0 else 0
        sensitivity = TP / (TP + FN) if (TP + FN) != 0 else 0
        specificity = TN / (TN + FP) if (TN + FP) != 0 else 0
        f1_or_dsc = 2 * TP / (2 * TP + FP + FN) if (2 * TP + FP + FN) != 0 else 0
        miou = TP / (TP + FP + FN) if (TP + FP + FN) != 0 else 0

        log_info = f'val epoch: {epoch}, loss: {np.mean(loss_list):.4f}, miou: {miou}, f1_or_dsc: {f1_or_dsc}, accuracy: {accuracy}, \
                    specificity: {specificity}, sensitivity: {sensitivity}, confusion_matrix: {confusion}'
        print(log_info)
        logger.info(log_info)

    else:
        log_info = f'val epoch: {epoch}, loss: {np.mean(loss_list):.4f}'
        print(log_info)
        logger.info(log_info)
    
    return np.mean(loss_list)




def test_one_epoch(test_loader,
                    model,
                    criterion,
                    logger,
                    config,
                    test_data_name=None):
    # switch to evaluate mode
    model.eval()
    preds = []
    gts = []
    loss_list = []
    with torch.no_grad():
        for i, data in enumerate(tqdm(test_loader)):
            img, msk = data
            img, msk = img.cuda(non_blocking=True).float(), msk.cuda(non_blocking=True).float()
            out = model(img)
            loss = criterion(out, msk)
            loss_list.append(loss.item())
            msk = msk.squeeze(1).cpu().detach().numpy()
            gts.append(msk)
            if type(out) is tuple:
                out = out[0]
            out = out.squeeze(1).cpu().detach().numpy()
            preds.append(out) 
            save_imgs(img, msk, out, i, config.work_dir + 'outputs/', config.datasets, config.threshold, test_data_name=test_data_name)

        preds = np.array(preds).reshape(-1)
        gts = np.array(gts).reshape(-1)

        y_pre = np.where(preds>=config.threshold, 1, 0)
        y_true = np.where(gts>=0.5, 1, 0)
        
        
        
        #print("y_true:", y_true[:5])  
        #print("y_pre:", y_pre[:5])

        confusion = confusion_matrix(y_true, y_pre)
        if confusion.shape != (2, 2):  
            print("混淆矩阵的形状不正确，应为 2x2。当前形状为:", confusion.shape)
        TN, FP, FN, TP = confusion[0,0], confusion[0,1], confusion[1,0], confusion[1,1] 
        
        accuracy = float(TN + TP) / float(np.sum(confusion)) if float(np.sum(confusion)) != 0 else 0
        sensitivity = float(TP) / float(TP + FN) if float(TP + FN) != 0 else 0
        specificity = float(TN) / float(TN + FP) if float(TN + FP) != 0 else 0
        f1_or_dsc = float(2 * TP) / float(2 * TP + FP + FN) if float(2 * TP + FP + FN) != 0 else 0
        miou = float(TP) / float(TP + FP + FN) if float(TP + FP + FN) != 0 else 0

        if test_data_name is not None:
            log_info = f'test_datasets_name: {test_data_name}'
            print(log_info)
            logger.info(log_info)
        log_info = f'test of best model, loss: {np.mean(loss_list):.4f},miou: {miou}, f1_or_dsc: {f1_or_dsc}, accuracy: {accuracy}, \
                specificity: {specificity}, sensitivity: {sensitivity}, confusion_matrix: {confusion}'
        print(log_info)
        logger.info(log_info)
        
        
        

    return np.mean(loss_list)

